import os
import redis
import json
import logging
from dotenv import load_dotenv

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("db_connection")

# 加载.env文件
load_dotenv()

def get_redis_connection():
    """
    从.env文件中获取Redis连接URL并建立连接
    
    Returns:
        redis.Redis: Redis连接实例
    """
    # 直接获取REDIS_URL_RENDER
    redis_url_render = os.getenv("REDIS_URL_RENDER")
    
    # 优先使用REDIS_URL_RENDER，如果不为空
    if redis_url_render and redis_url_render.strip():
        redis_url = redis_url_render
        logger.info(f"使用Redis连接: REDIS_URL_RENDER -> {redis_url}")
    else:
        # 如果REDIS_URL_RENDER为空，则使用CHOOSE_DB_URL_NAME指定的URL
        chosen_db_url_name = os.getenv("CHOOSE_DB_URL_NAME", "REDIS_URL_LOCAL")
        redis_url = os.getenv(chosen_db_url_name, "redis://localhost:6379")
        logger.info(f"使用Redis连接: {chosen_db_url_name} -> {redis_url}")
    
    # 创建Redis连接
    try:
        if "://" in redis_url:
            # 使用URI格式连接
            conn = redis.from_url(redis_url, decode_responses=True)
        else:
            # 假设是简单的host:port格式
            host, port = redis_url.split(":")
            conn = redis.Redis(host=host, port=int(port), decode_responses=True)
        
        # 测试连接
        conn.ping()
        logger.info("成功连接到Redis数据库")
        return conn
    except Exception as e:
        logger.error(f"连接Redis时出错: {str(e)}")
        # 返回一个默认连接作为备用
        logger.warning("使用本地默认连接作为备用")
        return redis.Redis(host='localhost', port=6379, decode_responses=True)

def init_redis_from_json(json_file_path=None):
    """
    从JSON文件初始化Redis数据库
    
    Args:
        json_file_path (str, optional): JSON文件路径，默认为"json_files/ocr_result.json"
    
    Returns:
        bool: 是否成功初始化
    """
    if json_file_path is None:
        json_file_path = "json_files/ocr_result.json"
    
    # 检查文件是否存在
    if not os.path.exists(json_file_path):
        logger.warning(f"初始化文件不存在: {json_file_path}")
        return False
    
    try:
        # 获取Redis连接
        r = get_redis_connection()
        
        # 检查是否已初始化
        if r.exists("db_initialized"):
            logger.info("数据库已初始化，跳过加载初始数据")
            return True
        
        # 读取JSON文件
        with open(json_file_path, 'r', encoding='utf-8') as f:
            company_data = json.load(f)
        
        loaded_count = 0
        
        # 存储数据到Redis
        for company_id, details in company_data.items():
            # 使用db:default前缀
            key = f"db:default:company:{company_id}"
            
            # 添加记录
            for field, value in details.items():
                r.hset(key, field, value)
            loaded_count += 1
        
        # 设置默认数据库名称
        r.sadd("database_list", "default")
        
        # 标记数据库已初始化
        r.set("db_initialized", "1")
        
        logger.info(f"数据库初始化完成，加载了 {loaded_count} 条记录")
        return True
    except Exception as e:
        logger.error(f"初始化数据库时出错: {str(e)}")
        return False

# 当文件直接运行时，执行初始化
if __name__ == "__main__":
    # 初始化数据库
    init_redis_from_json() 